Skip to content

Align quantized_kv_start defaults with the CLI - #1652

Open
Vlor999 wants to merge 1 commit into
ml-explore:mainfrom
Vlor999:fix/quantized-kv-start-default-v2
Open

Align quantized_kv_start defaults with the CLI#1652
Vlor999 wants to merge 1 commit into
ml-explore:mainfrom
Vlor999:fix/quantized-kv-start-default-v2

Conversation

@Vlor999

@Vlor999 Vlor999 commented Aug 3, 2026

Copy link
Copy Markdown

Proposed changes

Fixes #1651.

generate_step() and speculative_generate_step() defaulted quantized_kv_start
to 0, while the CLI already ships DEFAULT_QUANTIZED_KV_START = 5000. A caller
passing kv_bits= without quantized_kv_start therefore got different behaviour
depending on whether they came through the library or the CLI. This uses the
existing constant as the default in both functions and documents the trade-off
it encodes.

Why

mlx-community/Qwen3-0.6B-4bit, kv_bits=8, 8-token prompt, median of 3,
Apple M4 Pro (14-core), macOS 26.6:

new tokens start=0 (before) start=5000 (after) speed peak mem before peak mem after
64 325.4 tok/s 375.4 tok/s +15.4% 367.2 MB 367.7 MB
128 314.2 tok/s 382.3 tok/s +21.7% 367.2 MB 367.7 MB
256 336.3 tok/s 381.9 tok/s +13.6% 367.4 MB 421.7 MB
512 335.9 tok/s 368.1 tok/s +9.6% 371.2 MB 471.7 MB
1024 322.8 tok/s 341.9 tok/s +5.9% 408.2 MB 576.8 MB

The trade-off, stated plainly

This is not a free win. The new default is 6-22% faster to decode, but it
uses more memory once generations get long — 169 MB more at 1024 tokens. Below
roughly 256 generated tokens the memory is a wash (367 MB either way) and the
speed difference is the whole story; past that, the old default genuinely buys
memory back.

So the argument here is consistency, not superiority: the CLI already
encodes 5000 as the intended default via a shared constant, and the library
silently disagreeing with it is the part that looks like a bug. If you would
rather keep 0 and document the difference instead, say so and I will close
this in favour of a docs-only change — I raised both options in #1651.

This changes behaviour for existing callers who pass kv_bits= and rely on
the current 0. They can restore it explicitly with quantized_kv_start=0,
which the updated docstring now calls out.

Validation

python -m unittest tests.test_generate tests.test_prompt_cache tests.test_utils
# 60 passed
pre-commit run --files mlx_lm/generate.py tests/test_generate.py   # black, isort

Three tests added, pinning both function defaults to DEFAULT_QUANTIZED_KV_START
and asserting cache_prompt.py still shares the same value — so the two cannot
drift apart again silently.

I checked the tests are not vacuous by reverting generate_step's default to
0: the suite fails. With the change in place, all 60 pass.


Supersedes #1567, which was opened against an older main. Same change, rebuilt
on current main, with the measurements and tests added.

generate_step() and speculative_generate_step() defaulted quantized_kv_start
to 0, so a caller passing kv_bits= without it quantized the KV cache from the
first token. The CLI already ships DEFAULT_QUANTIZED_KV_START = 5000, so the
same arguments behaved differently depending on the entry point used.

Use the existing constant as the default in both functions and document the
trade-off it encodes.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

quantized_kv_start default disagrees between generate_step and the CLI

1 participant